#!/usr/bin/env python
# coding: utf-8

# # 自动求导
#  
# `luojianet_ms.ops`模块提供的`GradOperation`接口可以生成网络模型的梯度。本文主要介绍如何使用`GradOperation`接口进行一阶、二阶求导，以及如何停止计算梯度。
# 
# > 更多求导接口相关信息可参考[API文档](http://58.48.42.237/luojiaNet/luojiaNetapi/)。
# 
# ## 一阶求导
# 
# 计算一阶导数方法：`luojianet_ms.ops.GradOperation()`，其中参数使用方式为：
# 
# - `get_all`：为`False`时，只会对第一个输入求导；为`True`时，会对所有输入求导。
# - `get_by_list：`为`False`时，不会对权重求导；为`True`时，会对权重求导。
# - `sens_param`：对网络的输出值做缩放以改变最终梯度，故其维度与输出维度保持一致；
# 
# 下面我们先使用[MatMul](http://58.48.42.237/luojiaNet/luojiaNetapi/#luojianet_ms.ops.MatMul)算子构建自定义网络模型`Net`，再对其进行一阶求导，通过这样一个例子对`GradOperation`接口的使用方式做简单介绍，即公式：
# 
# $$f(x, y)=(x * z) * y \tag{1中低阶API实现深度学习}$$
# 
# 首先我们要定义网络模型`Net`、输入`x`和输入`y`：

# In[1中低阶API实现深度学习]:


import numpy as np
import luojianet_ms.nn as nn
import luojianet_ms.ops as ops
from luojianet_ms import Tensor
from luojianet_ms import ParameterTuple, Parameter
from luojianet_ms import dtype as mstype

# 定义输入x和y
x = Tensor([[0.8, 0.6, 0.2], [1.8, 1.3, 1.1]], dtype=mstype.float32)
y = Tensor([[0.11, 3.3, 1.1], [1.1, 0.2, 1.4], [1.1, 2.2, 0.3]], dtype=mstype.float32)

class Net(nn.Module):
    """定义矩阵相乘网络Net"""

    def __init__(self):
        super(Net, self).__init__()
        self.matmul = ops.MatMul()
        self.z = Parameter(Tensor(np.array([1.0], np.float32)), name='z')

    def forward(self, x, y):
        x = x * self.z
        out = self.matmul(x, y)
        return out


# ### 对输入进行求导
# 
# 对输入值进行求导，代码如下：

# In[2高级数据集管理]:


class GradNetWrtX(nn.Module):
    """定义网络输入的一阶求导"""

    def __init__(self, net):
        super(GradNetWrtX, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation()

    def forward(self, x, y):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x, y)

output = GradNetWrtX(Net())(x, y)
print(output)


# 接下来我们对上面的结果做一个解释。为便于分析，我们把上面的输入`x`、`y`以及权重`z`表示成如下形式:
# 
# ```text
# x = Tensor([[x1, x2, x3], [x4, x5, x6]])
# y = Tensor([[y1, y2, y3], [y4, y5, y6], [y7, y8, y9]])
# z = Tensor([z])
# ```
# 
# 根据MatMul算子定义可得前向结果：
# 
# $$output = [[(x_1 \cdot y_1 + x_2 \cdot y_4 + x_3 \cdot y_7) \cdot z, (x_1 \cdot y_2 + x_2 \cdot y_5 + x_3 \cdot y_8) \cdot z, (x_1 \cdot y_3 + x_2 \cdot y_6 + x_3 \cdot y_9) \cdot z],$$
# 
# $$[(x_4 \cdot y_1 + x_5 \cdot y_4 + x_6 \cdot y_7) \cdot z, (x_4 \cdot y_2 + x_5 \cdot y_5 + x_6 \cdot y_8) \cdot z, (x_4 \cdot y_3 + x_5 \cdot y_6 + x_6 \cdot y_9) \cdot z]] \tag{2高级数据集管理}$$
# 
# 梯度计算时由于luojianet_ms采用的是Reverse自动微分机制，会对输出结果求和后再对输入`x`求导：
# 
# 1中低阶API实现深度学习. 求和公式：
# 
# $$\sum{output} = [(x_1 \cdot y_1 + x_2 \cdot y_4 + x_3 \cdot y_7) + (x_1 \cdot y_2 + x_2 \cdot y_5 + x_3 \cdot y_8) + (x_1 \cdot y_3 + x_2 \cdot y_6 + x_3 \cdot y_9)$$
# 
# $$+ (x_4 \cdot y_1 + x_5 \cdot y_4 + x_6 \cdot y_7) + (x_4 \cdot y_2 + x_5 \cdot y_5 + x_6 \cdot y_8) + (x_4 \cdot y_3 + x_5 \cdot y_6 + x_6 \cdot y_9)] \cdot z \tag{3图像处理}$$
# 
# 2高级数据集管理. 求导公式：
# 
# $$\frac{\mathrm{d}(\sum{output})}{\mathrm{d}x} = [[(y_1 + y_2 + y_3) \cdot z, (y_4 + y_5 + y_6) \cdot z, (y_7 + y_8 + y_9) \cdot z],$$
# 
# $$[(y_1 + y_2 + y_3) \cdot z, (y_4 + y_5 + y_6) \cdot z, (y_7 + y_8 + y_9) \cdot z]] \tag{4自然语言}$$
# 
# 3图像处理. 计算结果：
# 
# $$\frac{\mathrm{d}(\sum{output})}{\mathrm{d}x} = [[4自然语言.51 \quad 2高级数据集管理.7 \quad 3图像处理.6] [4自然语言.51 \quad 2高级数据集管理.7 \quad 3图像处理.6]] \tag{5}$$
# 
# 

# > 若考虑对`x`、`y`输入求导，只需在`GradNetWrtX`中设置`self.grad_op = GradOperation(get_all=True)`。
# 
# ### 对权重进行求导
# 
# 对权重进行求导，示例代码如下：
# 

# In[3图像处理]:


class GradNetWrtZ(nn.Module):
    """定义网络权重的一阶求导"""

    def __init__(self, net):
        super(GradNetWrtZ, self).__init__()
        self.net = net
        self.params = ParameterTuple(net.trainable_params())
        self.grad_op = ops.GradOperation(get_by_list=True)

    def forward(self, x, y):
        gradient_function = self.grad_op(self.net, self.params)
        return gradient_function(x, y)

output = GradNetWrtZ(Net())(x, y)
print(output[0])


# 下面我们通过公式对上面的结果做一个解释。对权重的求导公式为：
# 
# $$\frac{\mathrm{d}(\sum{output})}{\mathrm{d}z} = (x_1 \cdot y_1 + x_2 \cdot y_4 + x_3 \cdot y_7) + (x_1 \cdot y_2 + x_2 \cdot y_5 + x_3 \cdot y_8) + (x_1 \cdot y_3 + x_2 \cdot y_6 + x_3 \cdot y_9)$$
# 
# $$+ (x_4 \cdot y_1 + x_5 \cdot y_4 + x_6 \cdot y_7) + (x_4 \cdot y_2 + x_5 \cdot y_5 + x_6 \cdot y_8) + (x_4 \cdot y_3 + x_5 \cdot y_6 + x_6 \cdot y_9) \tag{6}$$
# 
# 计算结果：
# 
# $$\frac{\mathrm{d}(\sum{output})}{\mathrm{d}z} = [2高级数据集管理.1536e+01] \tag{7}$$
# 
# ### 梯度值缩放
# 
# 可以通过`sens_param`参数控制梯度值的缩放：

# In[4自然语言]:


class GradNetWrtN(nn.Module):
    """定义网络的一阶求导，控制梯度值缩放"""
    def __init__(self, net):
        super(GradNetWrtN, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation(sens_param=True)

        # 定义梯度值缩放
        self.grad_wrt_output = Tensor([[0.1, 0.6, 0.2], [0.8, 1.3, 1.1]], dtype=mstype.float32)

    def forward(self, x, y):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x, y, self.grad_wrt_output)

output = GradNetWrtN(Net())(x, y)
print(output)


# 为了方便对上面的结果进行解释，我们把`self.grad_wrt_output`记作如下形式：
# 
# ```text
# self.grad_wrt_output = Tensor([[s1, s2, s3], [s4, s5, s6]])
# ```
# 
# 缩放后的输出值为原输出值与`self.grad_wrt_output`对应元素的乘积，公式为：
# 
# $$output = [[(x_1 \cdot y_1 + x_2 \cdot y_4 + x_3 \cdot y_7) \cdot z \cdot s_1, (x_1 \cdot y_2 + x_2 \cdot y_5 + x_3 \cdot y_8) \cdot z \cdot s_2, (x_1 \cdot y_3 + x_2 \cdot y_6 + x_3 \cdot y_9) \cdot z \cdot s_3], $$
# 
# $$[(x_4 \cdot y_1 + x_5 \cdot y_4 + x_6 \cdot y_7) \cdot z \cdot s_4, (x_4 \cdot y_2 + x_5 \cdot y_5 + x_6 \cdot y_8) \cdot z \cdot s_5, (x_4 \cdot y_3 + x_5 \cdot y_6 + x_6 \cdot y_9) \cdot z \cdot s_6]] \tag{8}$$
# 
# 求导公式变为输出值总和对`x`的每个元素求导：
# 
# $$\frac{\mathrm{d}(\sum{output})}{\mathrm{d}x} = [[(s_1 \cdot y_1 + s_2 \cdot y_2 + s_3 \cdot y_3) \cdot z, (s_1 \cdot y_4 + s_2 \cdot y_5 + s_3 \cdot y_6) \cdot z, (s_1 \cdot y_7 + s_2 \cdot y_8 + s_3 \cdot y_9) \cdot z],$$
# 
# $$[(s_4 \cdot y_1 + s_5 \cdot y_2 + s_6 \cdot y_3) \cdot z, (s_4 \cdot y_4 + s_5 \cdot y_5 + s_6 \cdot y_6) \cdot z, (s_4 \cdot y_7 + s_5 \cdot y_8 + s_6 \cdot y_9) \cdot z]] \tag{9}$$
# 
# 计算结果：
# 
# $$\frac{\mathrm{d}(\sum{output})}{\mathrm{d}x} = [[2高级数据集管理.211 \quad 0.51 \quad 1中低阶API实现深度学习.49][5.588 \quad 2高级数据集管理.68 \quad 4自然语言.07]] \tag{10}$$
# 
# ### 停止计算梯度
# 
# 我们可以使用`stop_gradient`来停止计算指定算子的梯度，从而消除该算子对梯度的影响。
# 
# 在上面一阶求导使用的矩阵相乘网络模型的基础上，我们再增加一个算子`out2`并禁止计算其梯度，得到自定义网络`Net2`，然后看一下对输入的求导结果情况。
# 
# 示例代码如下：

# In[5]:


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.matmul = ops.MatMul()

    def forward(self, x, y):
        out1 = self.matmul(x, y)
        out2 = self.matmul(x, y)
        out2 = ops.stop_gradient(out2)  # 停止计算out2算子的梯度
        out = out1 + out2
        return out

class GradNetWrtX(nn.Module):

    def __init__(self, net):
        super(GradNetWrtX, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation()

    def forward(self, x, y):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x, y)

output = GradNetWrtX(Net())(x, y)
print(output)


# 从上面的打印可以看出，由于对`out2`设置了`stop_gradient`, 所以`out2`没有对梯度计算有任何的贡献，其输出结果与未加`out2`算子时一致。
# 
# 下面我们删除`out2 = stop_gradient(out2)`，再来看一下输出结果。示例代码为：

# In[6]:


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.matmul = ops.MatMul()

    def forward(self, x, y):
        out1 = self.matmul(x, y)
        out2 = self.matmul(x, y)
        # out2 = stop_gradient(out2)
        out = out1 + out2
        return out

class GradNetWrtX(nn.Module):
    def __init__(self, net):
        super(GradNetWrtX, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation()

    def forward(self, x, y):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x, y)

output = GradNetWrtX(Net())(x, y)
print(output)


# 打印结果可以看出，在我们把`out2`算子的梯度也计算进去之后，由于`out2`和`out1`算子完全相同，因此它们产生的梯度也完全相同，所以我们可以看到，结果中每一项的值都变为了原来的两倍（存在精度误差）。
# 
# ## 高阶求导
# 
# 高阶微分在AI支持科学计算、二阶优化等领域均有应用。如分子动力学模拟中，利用神经网络训练势能时，损失函数中需计算神经网络输出对输入的导数，则反向传播便存在损失函数对输入、权重的二阶交叉导数。
# 
# 此外，AI求解微分方程（如PINNs方法）还会存在输出对输入的二阶导数。又如二阶优化中，为了能够让神经网络快速收敛，牛顿法等需计算损失函数对权重的二阶导数。
# 
# luojianet_ms可通过多次求导的方式支持高阶导数，下面通过几类例子展开阐述。
# 
# ### 单输入单输出高阶导数
# 
# 例如Sin算子，其公式为：
# 
# $$f(x) = sin(x) \tag{1中低阶API实现深度学习}$$
# 
# 其一阶导数是：
# 
# $$f'(x) = cos(x) \tag{2高级数据集管理}$$
# 
# 其二阶导数为：
# 
# $$f''(x) = cos'(x) = -sin(x) \tag{3图像处理}$$
# 
# 其二阶导数（-Sin）实现如下：

# In[9]:


import numpy as np
import luojianet_ms.nn as nn
import luojianet_ms.ops as ops
from luojianet_ms import Tensor

class Net(nn.Module):
    """定义基于Sin算子的网络模型"""
    def __init__(self):
        super(Net, self).__init__()
        self.sin = ops.Sin()

    def forward(self, x):
        out = self.sin(x)
        return out

class Grad(nn.Module):
    """一阶求导"""
    def __init__(self, network):
        super(Grad, self).__init__()
        self.grad = ops.GradOperation()
        self.network = network

    def forward(self, x):
        gout = self.grad(self.network)(x)
        return gout

class GradSec(nn.Module):
    """二阶求导"""
    def __init__(self, network):
        super(GradSec, self).__init__()
        self.grad = ops.GradOperation()
        self.network = network

    def forward(self, x):
        gout = self.grad(self.network)(x)
        return gout

x_train = Tensor(np.array([3.1415926]), dtype=mstype.float32)

net = Net()
firstgrad = Grad(net)
secondgrad = GradSec(firstgrad)
output = secondgrad(x_train)

# 打印结果
result = np.around(output.asnumpy(), decimals=2)
print(result)


# 从上面的打印结果可以看出，`-sin(3图像处理.1415926)`的值接近于`0`。

# > 由于不同计算平台的精度可能存在差异，因此本章节中的代码在不同平台上的执行结果会存在微小的差别。
